#!/usr/bin/python

"""
Binary Tree implementation and visualization
"""

import os
import sys 

pygraphvizOk = True
try:
    import pygraphviz
except ImportError:
      pygraphvizOk = False


__metaclass__ = type

class BinTree:
    def __init__(self, value = 0, parent = None, label='', left=None, right=None):
        self.value=value
        self.left=left
        self.right=right
        self.parent = parent
        self.count = 1
        if label == '':
            self.label = str(value)
        else:
            self.label = label

    def __del__(self):
        if self.right:
            del self.right
        if self.left:
            del self.left
        
    
    def buildFromArray(self, array):
        """
        Build a binary tree from linear array, s.t. root value is set to array[0];
        for nodes with value array[i], its left child has value array[i*2+1] and right child has array[i*2+2]
        """
        self.insertFromArray(array, 0)

    def insertFromArray(self, array, array_pos):
        """
        recursively set tree node value from array, 
        so that self get value array[array_pos] and conforms to rules for buildFromArray
        """
        array_len = len(array)

        self.value = array[array_pos]
        self.label = str(self.value)
        left_pos = array_pos * 2 + 1
        right_pos = left_pos + 1
        if right_pos < array_len:
            child = self.insertLeft()
            child.insertFromArray(array, left_pos)
            child = self.insertRight()
            child.insertFromArray(array, right_pos)
        elif left_pos < array_len:
            child = self.insertLeft()
            child.insertFromArray(array, left_pos)

    def insertLeft(self,value = 0):
        if not self.left:
            self.left=BinTree(value, self)
        
        child = self.left
        child.value = value
        child.label = str(value)
        return child

    def insertRight(self,value=0):
        if not self.right:
            self.right=BinTree()
        
        child = self.right
        child.value = value
        child.parent = self 
        child.label = str(value)
        return child

    def getNodeCount(self):
        """
        Number of nodes in the tree with self as root
        """
        count = 1
        if self.right:
            count += self.right.getNodeCount()
        if self.left:
            count += self.left.getNodeCount()
        return count
    
    def getTreeHeight(tree):
        """
        Return height of a tree, s.t. a single-node-tree has height 1
        """
        if tree.right == None:
            rightH = 0
        else:
            rightH = tree.right.getTreeHeight()

        if tree.left == None:
            leftH = 0
        else:
            leftH = tree.left.getTreeHeight()

        return max([leftH, rightH]) + 1
    

    def paintTree(self, painter):
        """
        Visualize the tree with any object featuring obj.paint(tree) interface 
        """
        painter.paint(self)
    
    def inOrderFuncTraverse(self, func):
        if self.left:
            self.left.inOrderFuncTraverse(func)
        func(self)
        if self.right:
            self.right.inOrderTraverse(func)

    def preOrderFuncTraverse(self, func):
        func(self)
        if self.left:
            self.left.preOrderFuncTraverse(func)
        if self.right:
            self.right.preOrderFuncTraverse(func)
 
    def postOrderFuncTraverse(self, func):
        if self.left:
            self.left.postOrderFuncTraverse(func)
        if self.right:
            self.right.postOrderFuncTraverse(func)
        func(self)


class binSearchTree(BinTree):
    pass


class BinTreePainterBase:
    """
    Base class of binary tree painters that visualize BinTree objects
    """
    def __clear__(self):
        self.scriptFile = None

    def __init__(self, imageName, scriptSuffix='', imageType='png'):
        self.scriptName = imageName + '.' +scriptSuffix
        self.scriptFile = None
        self.imageType = imageType
        self.imageName = imageName+'.'+imageType
        self.__clear__()


    def display(self):
        if not os.path.isfile(self.imageName):
            return
        if self.imageType == 'pdf':
            os.system('evince %s' % self.imageName)
        elif self.imageType == 'png':
            os.system('eog %s' % self.imageName)
    
    def paintNode(self, node):
        pass

    def paintBegin(self, tree):
        pass
    
    def paintEnd(self, tree):
        pass
 
    def paint(self, tree):
        self.paintBegin(tree)
        tree.preOrderFuncTraverse(self.paintNode)
        self.paintEnd(tree)


class BinTreeTikztreePainter(BinTreePainterBase):
    """
    Binary tree painter, which can visualize a BinTree Object with latex's tikz tree drawing functionality
    """
    def __init__(self, imageName='tikzTree', imageType='pdf', sep=20):
        super(BinTreeTikztreePainter, self).__init__(imageName, 'tex', imageType)
        self.sep = sep

    def paintBegin(self, tree):
        height = tree.getTreeHeight()
        self.scriptFile.write("\\documentclass[11pt,a4paper]{article}\n\n")
        self.scriptFile.write("\\uself.sepackage{tikz}\n\n")
        self.scriptFile.write("\\begin{document}\n\n")
        self.scriptFile.write("\\tikzset{ box/.style={\n")
        self.scriptFile.write("rectangle, minimum width=.001\\textwidth, very thick, draw=gray!50!black!50, text centered,\n")
        self.scriptFile.write("top color=white, bottom color=gray!50!black!20},\n")
        self.scriptFile.write("every node/.style={text centered, font=\\footnotesize, text width=0.03\\textwidth}\n")
        self.scriptFile.write("}\n\n")
        self.scriptFile.write("\\begin{tikzpicture}[")
        while height > 1:
            self.scriptFile.write("level %d/.style={sibling distance=%dpt},\n"%(height, self.sep))
            height -= 1
            self.sep *= 2

        self.scriptFile.write("level 1/.style={sibling distance=%dpt}]\n\n\\"%self.sep)

    def paintEnd(self):
        self.scriptFile.write(";\n\\end{tikzpicture}\n")
        self.scriptFile.write("\\end{document}\n")
    
    def doPaint(self, node):
        assert node
        self.scriptFile.write("node[box] {%s}\n" % node.label)
        if node.left:
            self.scriptFile.write("child {\n")
            self.doPaint(node.left)
            self.scriptFile.write("}\n")
        if node.right:
            self.scriptFile.write("child {\n")
            self.doPaint(node.right)
            self.scriptFile.write("}\n")
    

    def __runTex__(self):
        if self.imageType == 'pdf':
            os.system('pdflatex %s' % self.scriptName)
            os.system('rm -f %s' % self.scriptName)
        else:
            print 'Formats other than PDF not supported' 

    def paint(self, tree):
        with open(self.scriptName, 'w') as self.scriptFile:
            self.paintBegin(tree)
            if tree:
                self.doPaint(tree)
            self.paintEnd()
            self.scriptFile.close()
            self.__runTex__()
            self.__clear__()

class BinTreeGraphvizPainter(BinTreePainterBase):
    """
    Binary tree painter, which visualizes a BinTree Object with Graphviz's dot language
    """
    def __init__(self, imageName='graphvizTree', imageType='png'):
        super(BinTreeGraphvizPainter, self).__init__(imageName, 'dot', imageType)
    
    def paintNode(self, node):
        if not node: return
        self.scriptFile.write('%d[label="%s"];\n' % (id(node), node.label))
        if node.left and node.right:
            self.scriptFile.write('%d->{%d; %d};\n' % (id(node), id(node.left), id(node.right)))
        elif node.left:
            self.scriptFile.write('%d->%d;\n' % (id(node), id(node.left)))
        elif node.right:
            self.scriptFile.write('%d->%d;\n' % (id(node), id(node.right)))

    def paintBegin(self, tree):
        if not tree: return
        self.scriptFile = open(self.scriptName, 'w')
        self.scriptFile.write('digraph g {\n')
        self.scriptFile.write(' node[shape=record];\n')
    
    def paintEnd(self, tree):
        self.scriptFile.write('}')
        self.scriptFile.close()
        os.system('dot -T%s -o %s %s' % (self.imageType, self.imageName, self.scriptName))
        self.__clear__()

 
class BinTreePinpointPainter(BinTreePainterBase):
    """
    Base to some binary tree painters, 
    which visualize a BinTree Object with precise tree node X-Y positioning
    """
    def __clear__(self):
        self.curXpos = 0 
        self.curYpos = 0 
        self.graphMap = []

    def __init__(self, imageName='absPosTree', scriptSuffix='dot', imageType='png'):
        super(BinTreePinpointPainter, self).__init__(imageName, scriptSuffix, imageType)
        self.__clear__()
    
    
    def doPinPoint(self, node):
        assert node
        if node.left:
            self.curYpos -= 1
            self.doPinPoint(node.left)
            self.curYpos += 1
        self.graphMap[self.curYpos].append((node, self.curXpos))
        if node.right:
            self.curXpos += 1
            self.curYpos -= 1
            self.doPinPoint(node.right)
            self.curYpos += 1

    def pinPoint(self, tree):
        if not tree: return
        self.__clear__()
        height = tree.getTreeHeight()
        for i in range(0, height):
            self.graphMap.append([])
        self.curYpos = height - 1
        self.doPinPoint(tree)
    
    def paintBegin(self, tree):
        pass

    def paintEnd(self, tree):
        pass

    def paintNode(self, node, xpos, ypos):
        pass


    def paint(self, tree):
        self.pinPoint(tree)
        height = tree.getTreeHeight()
        self.paintBegin(tree)
        while height > 0:
            height -= 1
            for mapNode in self.graphMap[height]:
                self.paintNode(mapNode[0], mapNode[1], height)
        self.paintEnd(tree)
        self.__clear__()

if pygraphvizOk:
    #Only when python-pygraphviz is installed
    class BinTreePygraphvizPainter(BinTreePinpointPainter):
        """
        Binary tree painter, which visualizes a BinTree Object with Pygraphviz module 
        """
        def __init__(self, imageName='pygraphvizTree', imageType='png'):
            super(BinTreePygraphvizPainter, self).__init__(imageName, 'dot', imageType)

        def paintBegin(self, tree):
            self.graph = pygraphviz.AGraph(directed=True)

        def paintEnd(self, tree):
            self.graph.layout()
            self.graph.draw(self.imageName, self.imageType)
            self.graph = None

        def paintNode(self, node, xpos, ypos):
            self.graph.add_node(id(node), label=node.label)
            graphNode = self.graph.get_node(id(node)) 
            graphNode.attr['pos'] = '%d,%d!'%(xpos, ypos*1.5)
            if node.left:
                self.graph.add_node(id(node.left), label=node.left.label)
                self.graph.add_edge(id(node), id(node.left))
            if node.right:
                self.graph.add_node(id(node.right), label=node.right.label)
                self.graph.add_edge(id(node), id(node.right))


class BinTreeTextPainter(BinTreePinpointPainter):
    """
    Binary tree text painter. Paints on stdout. 
    """
    def __init__(self, imageName='textTree', imageType='txt'):
        super(BinTreeTextPainter, self).__init__(imageName, 'txt', imageType)

    def paintBegin(self, tree):
        self.curXpos = 0
        self.curYpos = tree.getTreeHeight() - 1

    def paintEnd(self, tree):
        print '\n'

    def paintNode(self, node, xpos, ypos):
        if self.curYpos == ypos + 1:
            print '\n'
            self.curYpos -= 1
            self.curXpos = 0
        
        if self.curYpos != ypos: return

        assert self.curXpos <= xpos
        while self.curXpos < xpos:
            self.curXpos += 1
            print '    ',
        print '%4s'%node.label,
        self.curXpos += 1
        

class BinTreeTikzPainter(BinTreePinpointPainter):
    """
    Binary tree painter, which visualizes a BinTree Object with latex's tikz node drawing functionality
    """

    def __init__(self, imageName='tikznodeTree', imageType='pdf'):
        super(BinTreeTikzPainter, self).__init__(imageName, 'tex', imageType)

    def __runTex__(self):
        if self.imageType == 'pdf':
            os.system('pdflatex %s' % self.scriptName)
            os.system('rm -f %s' % self.scriptName)
        else:
            print 'Formats other than PDF not supported' 

    def paintBegin(self, tree):
        self.scriptFile = open(self.scriptName, 'w')
        self.scriptFile.write("\\documentclass[11pt,a4paper]{article}\n\n")
        self.scriptFile.write("\\usepackage{tikz}\n\n")
        self.scriptFile.write("\\begin{document}\n\n")
        self.scriptFile.write("\\tikzset{ box/.style={\n")
        self.scriptFile.write("rectangle, minimum width=.001\\textwidth, very thick, draw=gray!50!black!50, text centered,\n")
        self.scriptFile.write("top color=white, bottom color=gray!50!black!20},\n")
        self.scriptFile.write("every node/.style={text centered, font=\\footnotesize, text width=0.03\\textwidth}\n")
        self.scriptFile.write("}\n\n")
        self.scriptFile.write("\\begin{tikzpicture}[every path/.style={draw, thick, ->}]")

    def paintEnd(self, tree):
        self.scriptFile.write("\n\\end{tikzpicture}\n")
        self.scriptFile.write("\\end{document}\n")
        self.scriptFile.close()
        self.__runTex__()

    def paintNode(self, node, xpos, ypos):
        self.scriptFile.write("\\node[box] at(%f, %f) (%s) {%s};\n\n" % (xpos, ypos, id(node), node.label))
        if node.parent:
            self.scriptFile.write("\path(%s) -- (%s);\n\n" % (id(node.parent), id(node)))

 
if __name__ == '__main__':

    import unittest

    class painterTest(unittest.TestCase):

        def testPainter(self):
            minv = -99
            maxv = 99
            x = [maxv]
            from random import randint
            n = randint(1, 19)
            for i in range(0, n):
                x.append(randint(minv+1, maxv))
            tree = binSearchTree()
            tree.buildFromArray(x[1:n+1])
            txtPainter = BinTreeTextPainter()
            tree.paintTree(txtPainter)
            if pygraphvizOk:
                painter = BinTreePygraphvizPainter()
                tree.paintTree(painter)
                painter.display()


    unittest.main()








